from typing import Sequence, Any

import numpy as np

import quimb as qu
import quimb.tensor as qtn

from quimb.tensor.tensor_1d import TensorNetwork1DOperator
from quimb.tensor import (
	Tensor, TensorNetwork, MPS_product_state
)
from quimb.core import make_immutable
from tqdm import trange

from states import basic_ptm_vec

sqrt2_I_o1 = basic_ptm_vec(0, 4).reshape(-1) * qu.sqrt(2)
make_immutable(sqrt2_I_o1)

def QPT_data_gen(
		N_shot: int,
		tn_1d_op: TensorNetwork,
		rhoin_map: dict[str, Any],
		meas_map: dict[str, dict],
		rhoin_setting: Sequence[str],
		meas_setting: Sequence[str],

		N_qubit: int = None,
		upper_ind_id: str = "k{}",
		lower_ind_id: str = "b{}",
		site_tag_id: str = "I{}",
		verbose: bool = True,
		**kwargs
) -> list[str]:
	"""
	N_shot: int
		The number of sampling times
	tn_1d_op: TensorNetwork1DOperator
		1d-like TensorNetwork, with `upper_ind_id`, `lower_ind_id`, `site_tag_id` and `nsites` property
	rhoin_map: dict
		A dictionary for possible input density operators in PTM form.
		*Key* corresponds to operator name, and *Val* corresponds to its PTM form.
		e.g.
		``rhoin_map = {
			'x+': np.array([1, 1, 0, 0]) / 2**0.5,
			'x-': np.array([1,-1, 0, 0]) / 2**0.5,
			'y+': np.array([1, 0, 1, 0]) / 2**0.5,
			'y-': np.array([1, 0,-1, 0]) / 2**0.5,
			'z+': np.array([1, 0, 0, 1]) / 2**0.5,
			'z-': np.array([1, 0, 0,-1]) / 2**0.5,
		}``
	meas_map: dict[str, dict]
		A dictionary for possible measurement operators in PTM form.
		*Key* corresponds to the Observable name, and *Val* corresponds to its possible effects. Furthermore,
		*Val* is a dictionary with Effect(Probability Operator) name as key and corresponding PTM form as val.
		e.g.
		``meas_map = {
			'x': {
					'x+': np.array([1, 1, 0, 0]) / 2**0.5,
					'x-': np.array([1,-1, 0, 0]) / 2**0.5
				},
			'y': {
					'y+': np.array([1, 0, 1, 0]) / 2**0.5,
					'y-': np.array([1, 0,-1, 0]) / 2**0.5
				},
			'z': {
					'z+': np.array([1, 0, 0, 1]) / 2**0.5,
					'z-': np.array([1, 0, 0,-1]) / 2**0.5
				},
		}``
	rhoin_setting: Sequence[str]
		The initial setting of input density PTM matrix on each qubit. Its values should be the key of ``rhoin_map``.
		e.g. ``rhoin_setting = ['x+', 'y-', 'z+', 'x-']`` for 4 qubits
	meas_setting: Sequence[str]
		The initial setting of measurement PTM Observable on each qubit. Its values should be the key of ``meas_map``.
		e.g. ``meas_setting = ['z', 'z', 'x', 'y']`` for 4 qubits
	N_qubit: int, optional
		The number of qubit. If is None, will be obtained from `tn_1d_op.nsites`.
	verbose: bool
		whether to show sampling process bar and probability.
	"""
	if isinstance(tn_1d_op, TensorNetwork1DOperator):
		upper_ind_id = tn_1d_op.upper_ind_id
		lower_ind_id = tn_1d_op.lower_ind_id
		site_tag_id = tn_1d_op.site_tag_id
		N_qubit = tn_1d_op.nsites
	if N_qubit is not None:
		assert len(rhoin_setting) == len(meas_setting) == N_qubit, \
			"The length of `rhoin_setting` and `meas_setting` should be the same as `N_qubit`"
	else:
		N_qubit = len(rhoin_setting)
		assert len(rhoin_setting) == len(meas_setting), \
			"The length of `rhoin_setting` and `meas_setting` should be the same"

	# MPS of input density operators in PTM form
	rhoi_mps = MPS_product_state(tuple(map(rhoin_map.get, rhoin_setting)), site_ind_id=lower_ind_id)
	rhoo_mps = rhoi_mps | tn_1d_op
	qtn.tensor_network_1d_compress(rhoo_mps, method='direct', inplace=True)

	pbar = trange(N_shot, desc=kwargs.get("desc", "Sampling"), leave=True) if verbose else range(N_shot)
	"""Sample `N_shot` times"""
	for _ in pbar:
		meas_mps = TensorNetwork(
			[Tensor(sqrt2_I_o1, inds=(upper_ind_id.format(_),), tags=[site_tag_id.format(_)]) for _ in
			 range(N_qubit)], virtual=True)
		out, prob = [], 1.0
		for j, jth_meas_setting in enumerate(meas_setting):
			meas_mps.pop_tensor(j)
			upper_ind, site_tag = upper_ind_id.format(j), site_tag_id.format(j)

			# Effects (Probability Operators) of the Observable of jth qubit
			effects = meas_map[jth_meas_setting]
			# Calculating all probabilities of the Observable of jth qubit
			probs = []
			for k, effect_k in enumerate(effects.values()):
				meas_mps_ = meas_mps | Tensor(effect_k, inds=(upper_ind,), tags=[site_tag])
				prob_k = np.real((rhoo_mps | meas_mps_).contract())
				probs.append(prob_k)
			# Monte-Carlo sample
			mc_bar = np.random.uniform(0, np.sum(probs))
			jth_out_ind = np.searchsorted(np.cumsum(probs), mc_bar)
			jth_out = tuple(effects.keys())[jth_out_ind]

			prob = probs[jth_out_ind]

			meas_mps |= Tensor(effects[jth_out], inds=(upper_ind,), tags=[site_tag])
			out.append(jth_out)
		yield out
		if verbose:
			pbar.set_postfix(prob=prob)
